import json


def readSubtitles(path_file):
    # 定义四个列表
    list1 = []
    list2 = []
    list3 = []
    list4 = []
    # 打开并读取txt文件
    with open(path_file, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    # 按每4行进行循环处理
    try:
        for i in range(0, len(lines), 4):
            # 分别将每四行的数据添加到对应的列表中
            list1.append(lines[i].strip())  # 第一行数据
            list2.append(lines[i + 1].strip())  # 第二行数据
            list3.append(lines[i + 2].strip())  # 第三行数据
            list4.append('')
    except:
        pass
        # list4.append(lines[i + 3].strip())  # 第四行数据
    # 输出结果
    # print("list1:", list1)
    # print("list2:", list2)
    # print("list3:", list3)
    # print("list4:", list4)
    return list1, list2, list3, list4


def writeSubtitles(path_file, list1, list2, list3, list4):
    # 确保四个列表长度一致
    assert len(list1) == len(list2) == len(list3) == len(list4)
    # 写入txt文件
    with open(path_file, 'w', encoding='utf-8') as f:
        for i in range(len(list1)):
            # 每次循环写入四行
            f.write(f"{list1[i]}\n")
            f.write(f"{list2[i]}\n")
            f.write(f"{list3[i]}\n")
            f.write(f"{list4[i]}\n")
    # 文件已写入完毕，可以查看example.txt文件内容



def generate_ollama_modelfile_message(path_file_ch, path_file_en, path_json):
    # 读取srt
    ch_list1, ch_list2, ch_list3, ch_list4 = readSubtitles(path_file_ch)
    en_list1, en_list2, en_list3, en_list4 = readSubtitles(path_file_en)
    # 解析&拼接
    json_list = []
    for i in range(len(ch_list2)):
        json_list.append({
            "no": ch_list1[i],
            "chinese_time": ch_list2[i],
            "chinese_subtitle": ch_list3[i],
            "english_time": en_list2[i+3],
            "english_subtitle": en_list3[i+3],
        })
    # 写文件
    with open(path_json,"w",encoding="utf8") as f:
        for d in json_list:
            f.write(json.dumps(d,ensure_ascii=False)+"\n")
        f.close()
    print("------------------------------")




def generate_colab_json(path_file_ch, path_file_en, path_json):
    # 读取srt
    ch_list1, ch_list2, ch_list3, ch_list4 = readSubtitles(path_file_ch)
    en_list1, en_list2, en_list3, en_list4 = readSubtitles(path_file_en)
    # 解析&拼接
    json_list = []
    for i in range(len(ch_list2)):
        json_list.append({
            "instruction": ch_list3[i],
            "input": "",
            "output": en_list3[i+3],
        })
    # 写文件
    with open(path_json, "w", encoding="utf8") as f:
        for d in json_list:
            f.write(json.dumps(d, ensure_ascii=False) + "\n")
        f.close()
    print("------------------------------")




if __name__ == '__main__':
    # testHelsinkiZh2En()
    # testTranslate()
    # subtitlesParse("./files/subtitles.txt", "./files/subtitles_en.txt")
    path_file_ch = '../files/cfd25t中文.srt'
    path_file_en = '../files/cfd25t英文.srt'
    path_json = '../files/cfd25t.json'
    generate_ollama_modelfile_message(path_file_ch, path_file_en, path_json)
    generate_colab_json(path_file_ch, path_file_en, '../files/cfd25t_colab.json')